#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/rand.h>
#include <openssl/err.h>
#include <openssl/ec.h>

// 默认 SM2 用户 ID
const unsigned char SM2_DEFAULT_USERID[] = "1234567812345678";
const size_t SM2_DEFAULT_USERID_LEN = 16;

// 错误处理函数
void handleErrors(void) {
    fprintf(stderr, "OpenSSL Error:\n");
    ERR_print_errors_fp(stderr);
    abort();
}

// 从 PEM 文件加载公钥或私钥
EVP_PKEY *load_pkey(const char *filename, int is_private) {
    BIO *bio = BIO_new_file(filename, "r");
    if (!bio) {
        perror("BIO_new_file");
        return NULL;
    }

    EVP_PKEY *pkey = NULL;
    if (is_private) {
        pkey = PEM_read_bio_PrivateKey(bio, NULL, NULL, NULL);
    } else {
        pkey = PEM_read_bio_PUBKEY(bio, NULL, NULL, NULL);
    }

    BIO_free(bio);
    if (!pkey) {
        fprintf(stderr, "Error loading key from %s\n", filename);
        handleErrors();
    }
    return pkey;
}

/**
 * @brief 从文件读取数据
 * @param filename 文件名
 * @param data (输出) 指向分配的内存
 * @param len (输出) 读取的字节数
 * @return 1 for success, 0 for failure
 */
int read_from_file(const char *filename, unsigned char **data, size_t *len) {
    FILE *fp = fopen(filename, "rb"); // "rb" - read binary
    if (!fp) {
        perror(filename);
        return 0;
    }
    
    fseek(fp, 0, SEEK_END);
    *len = ftell(fp);
    fseek(fp, 0, SEEK_SET);

    if (*len == 0) {
        fprintf(stderr, "File %s is empty.\n", filename);
        fclose(fp);
        return 0;
    }

    *data = malloc(*len);
    if (!*data) {
        perror("malloc read_from_file");
        fclose(fp);
        return 0;
    }

    if (fread(*data, 1, *len, fp) != *len) {
        fprintf(stderr, "Error reading from %s\n", filename);
        free(*data);
        *data = NULL;
        fclose(fp);
        return 0;
    }

    fclose(fp);
    printf("Bob: Read %zu bytes from %s\n", *len, filename);
    return 1;
}

/**
 * @brief Bob 的接收操作 (函数原型与之前相同)
 */
int bob_receive(EVP_PKEY *Pka, EVP_PKEY *Skb,
                unsigned char *C, size_t C_len,
                unsigned char *KC, size_t KC_len,
                unsigned char *S1, size_t S1_len) {
    
    printf("--- Bob 开始接收 ---\n");

    // 1. Sm2Very(Pka, C, S1)
    EVP_MD_CTX *ctx_md_vfy = EVP_MD_CTX_new();
    EVP_PKEY_CTX *pctx_vfy = NULL;
    if (!ctx_md_vfy) handleErrors();

    if (EVP_DigestVerifyInit(ctx_md_vfy, &pctx_vfy, EVP_sm3(), NULL, Pka) <= 0) handleErrors();
    if (EVP_PKEY_CTX_set1_id(pctx_vfy, SM2_DEFAULT_USERID, SM2_DEFAULT_USERID_LEN) <= 0) handleErrors();
    if (EVP_DigestVerifyUpdate(ctx_md_vfy, C, C_len) <= 0) handleErrors();
    
    int verify_status = EVP_DigestVerifyFinal(ctx_md_vfy, S1, S1_len);
    EVP_MD_CTX_free(ctx_md_vfy);

    if (verify_status == 1) {
        printf("Bob: Sm2Very(Pka, C, S1) 成功. 签名有效。\n");
    } else {
        fprintf(stderr, "Bob: Sm2Very(Pka, C, S1) 失败! 签名无效。\n");
        return 0;
    }

    // 2. Sm2Dec(Skb, KC) = k
    EVP_PKEY_CTX *ctx_sm2_dec = EVP_PKEY_CTX_new_from_pkey(NULL, Skb, NULL);
    if (!ctx_sm2_dec) handleErrors();
    if (EVP_PKEY_decrypt_init(ctx_sm2_dec) <= 0) handleErrors();

    unsigned char k[16];
    size_t k_len = sizeof(k);

    if (EVP_PKEY_decrypt(ctx_sm2_dec, k, &k_len, KC, KC_len) <= 0) {
        fprintf(stderr, "Bob: 解密对称密钥 k 失败。\n");
        handleErrors();
        EVP_PKEY_CTX_free(ctx_sm2_dec);
        return 0;
    }
    
    if (k_len != 16) {
         fprintf(stderr, "Bob: 解密出的密钥长度 (%zu) 不是 16 字节。\n", k_len);
         EVP_PKEY_CTX_free(ctx_sm2_dec);
         return 0;
    }

    printf("Bob: Sm2Dec(Skb, KC) = k 成功, 恢复 SM4 密钥 k\n");
    EVP_PKEY_CTX_free(ctx_sm2_dec);

    // 3. Sm4Dec(k, C) = P
    unsigned char *iv = C;
    unsigned char *C_data = C + 16;
    size_t C_data_len = C_len - 16;

    if (C_data_len <= 0) {
        fprintf(stderr, "Bob: 密文 C 长度异常 (<= 16)。\n");
        return 0;
    }

    EVP_CIPHER_CTX *ctx_sm4_dec = EVP_CIPHER_CTX_new();
    if (!ctx_sm4_dec) handleErrors();

    unsigned char *P_dec = malloc(C_data_len + 16); 
    if (!P_dec) {
        perror("malloc P_dec");
        return 0;
    }
    int out_len_dec = 0, final_len_dec = 0;

    if (EVP_DecryptInit_ex(ctx_sm4_dec, EVP_sm4_cbc(), NULL, k, iv) != 1) handleErrors();
    if (EVP_DecryptUpdate(ctx_sm4_dec, P_dec, &out_len_dec, C_data, C_data_len) != 1) handleErrors();
    if (EVP_DecryptFinal_ex(ctx_sm4_dec, P_dec + out_len_dec, &final_len_dec) != 1) {
        fprintf(stderr, "Bob: Sm4Dec 失败 (可能是密钥k或IV错误, 或padding损坏)。\n");
        handleErrors();
        free(P_dec);
        EVP_CIPHER_CTX_free(ctx_sm4_dec);
        return 0;
    }

    size_t P_dec_len = out_len_dec + final_len_dec;
    P_dec[P_dec_len] = '\0'; // 添加 null 终止符

    printf("Bob: Sm4Dec(k, C) = P 成功。\n");
    printf("***********************************\n");
    printf("Bob: 恢复的明文 P = '%s'\n", P_dec);
    printf("***********************************\n");
    printf("--- Bob 步骤完成 ---\n");

    free(P_dec);
    EVP_CIPHER_CTX_free(ctx_sm4_dec);
    return 1;
}

int main() {
    // --- 密钥加载 (Bob 只需要 Alice 的公钥 和 自己的私钥) ---
    EVP_PKEY *pkA = load_pkey("pkA.pem", 0);
    EVP_PKEY *skB = load_pkey("skB.pem", 1);

    if (!pkA || !skB) {
        fprintf(stderr, "Bob: 加载密钥失败 (需要 pkA.pem, skB.pem)\n");
        return 1;
    }

    // --- 从文件读取数字信封 ---
    unsigned char *C = NULL;
    size_t C_len = 0;
    unsigned char *KC = NULL;
    size_t KC_len = 0;
    unsigned char *S1 = NULL;
    size_t S1_len = 0;

    printf("Bob: 正在尝试读取数字信封文件...\n");
    if (!read_from_file("envelope.C", &C, &C_len) ||
        !read_from_file("envelope.KC", &KC, &KC_len) ||
        !read_from_file("envelope.S1", &S1, &S1_len)) {
        
        fprintf(stderr, "Bob: 读取信封文件失败。\n");
        fprintf(stderr, "请确保 'envelope.C', 'envelope.KC', 'envelope.S1' 文件存在。\n");
        
        EVP_PKEY_free(pkA);
        EVP_PKEY_free(skB);
        free(C); free(KC); free(S1);
        return 1;
    }
    printf("Bob: 信封文件加载完毕。\n\n");


    // --- Bob 执行接收操作 ---
    if (!bob_receive(pkA, skB, C, C_len, KC, KC_len, S1, S1_len)) {
        fprintf(stderr, "\nBob: 协议失败! 无法恢复明文。\n");
    } else {
        printf("\n>>> Bob: 协议成功完成! <<<\n");
    }

    // --- 清理 ---
    EVP_PKEY_free(pkA);
    EVP_PKEY_free(skB);
    free(C);
    free(KC);
    free(S1);
    
    return 0;
}
